#include "sprite_renderer.h"
#include <GL/glew.h>
#include "engine/utils/file_utils.h"

namespace kanon
{
    SpriteBatch::~SpriteBatch()
    {
        glDeleteVertexArrays(1, &vao);
        glDeleteBuffers(1, &vbo);
        glDeleteBuffers(1, &ibo);
    }

    SpriteBatch::SpriteBatch(uint32_t max_sprites)
    {
        max_vertices = max_sprites * 4;
        max_indices = max_sprites * 6;
        num_vertices = 0;
        num_indices = 0;
        max_texture_slots = 16;
        vertex_base = nullptr;
        index_base = nullptr;

        // create vao and vbo
        glGenVertexArrays(1, &vao);
        glBindVertexArray(vao);

        glGenBuffers(1, &vbo);
        glBindBuffer(GL_ARRAY_BUFFER, vbo);
        glBufferData(GL_ARRAY_BUFFER, max_vertices * sizeof(SpriteVertex), nullptr, GL_DYNAMIC_DRAW);

        // set vertex attributes
        glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, sizeof(SpriteVertex), (void*)offsetof(SpriteVertex, pos));
        glEnableVertexAttribArray(0);
        glVertexAttribPointer(1, 4, GL_FLOAT, GL_FALSE, sizeof(SpriteVertex), (void*)offsetof(SpriteVertex, color));
        glEnableVertexAttribArray(1);
        glVertexAttribPointer(2, 2, GL_FLOAT, GL_FALSE, sizeof(SpriteVertex), (void*)offsetof(SpriteVertex, uv));
        glEnableVertexAttribArray(2);
        glVertexAttribPointer(3, 1, GL_FLOAT, GL_FALSE, sizeof(SpriteVertex), (void*)offsetof(SpriteVertex, texture_slot));
        glEnableVertexAttribArray(3);

        glGenBuffers(1, &ibo);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ibo);
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, max_indices * sizeof(uint32_t), nullptr, GL_DYNAMIC_DRAW);
    }

    void SpriteBatch::Begin()
    {
        // map buffer
        glBindVertexArray(vao);
        glBindBuffer(GL_ARRAY_BUFFER, vbo);
        vertex_base = (SpriteVertex*)glMapBuffer(GL_ARRAY_BUFFER, GL_WRITE_ONLY);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ibo);
        index_base = (uint32_t*)glMapBuffer(GL_ELEMENT_ARRAY_BUFFER, GL_WRITE_ONLY);
        num_vertices = 0;
        num_indices = 0;

        next_texture_slot = 0;
        texture_id2slot.clear();
    }

    void SpriteBatch::End()
    {
        glUnmapBuffer(GL_ARRAY_BUFFER);
        vertex_base = nullptr;
        glUnmapBuffer(GL_ELEMENT_ARRAY_BUFFER);
        index_base = nullptr;
    }

    bool SpriteBatch::Submit(const SpriteQuad& sprite, const glm::mat4& transform)
    {
        if (num_indices + sprite.num_indices() > max_indices)
            return false;
        if (num_vertices + sprite.num_vertices() > max_vertices)
            return false;

        float texture_slot = -1.0f;
        if (sprite.texture.get())
        {
            uint32_t tid = sprite.texture->getId();
            auto it = texture_id2slot.find(tid);
            if (it == texture_id2slot.end())
            {
                texture_id2slot[tid] = next_texture_slot;
                next_texture_slot += 1;

                if (next_texture_slot > max_texture_slots)
                    return false;
            }
            texture_slot = (float)texture_id2slot[tid];
        }

        // copy vertex data
        SpriteVertex* data = vertex_base + num_vertices;
        data[0].pos = { sprite.pos.x, sprite.pos.y };
        data[0].color = sprite.color;
        data[0].uv = glm::vec2(0.0f, 0.0f);
        data[0].texture_slot = texture_slot;

        data[1].pos = { sprite.pos.x, sprite.pos.y + sprite.size.y };
        data[1].color = sprite.color;
        data[1].uv = glm::vec2(0.0f, 1.0f);
        data[1].texture_slot = texture_slot;

        data[2].pos = { sprite.pos.x + sprite.size.x, sprite.pos.y + sprite.size.y };
        data[2].color = sprite.color;
        data[2].uv = glm::vec2(1.0f, 1.0f);
        data[2].texture_slot = texture_slot;

        data[3].pos = { sprite.pos.x + sprite.size.x, sprite.pos.y };
        data[3].color = sprite.color;
        data[3].uv = glm::vec2(1.0f, 0.0f);
        data[3].texture_slot = texture_slot;

        // copy index data
        uint32_t* index_data = index_base + num_indices;
        index_data[0] = num_vertices + 0;
        index_data[1] = num_vertices + 1;
        index_data[2] = num_vertices + 2;
        index_data[3] = num_vertices + 0;
        index_data[4] = num_vertices + 2;
        index_data[5] = num_vertices + 3;

        num_vertices += 4;
        num_indices += 6;
        return true;
    }


    SpriteRenderer::SpriteRenderer()
    {
        std::string vertex_shader_path = FileUtils::getShaderPath("sprite_renderer.vs");
        std::string fragment_shader_path = FileUtils::getShaderPath("sprite_renderer.fs");
        shader = Shader::FromShaderFiles(vertex_shader_path, fragment_shader_path);
        batch = std::make_shared<SpriteBatch>(MAX_SPRITES);
    }

    bool SpriteRenderer::Submit(const SpriteQuad& sprite)
    {
        glm::mat4 t(1.0f);
        return this->Submit(sprite, t);
    }

    bool SpriteRenderer::Submit(const SpriteQuad& sprite, const glm::mat4& transform)
    {
        bool success = batch->Submit(sprite, transform);
        if (!success)
        {
            this->Flush();
            this->Begin();
            batch->Submit(sprite, transform);
        }
        return true;
    }

    void SpriteRenderer::Begin()
    {
        batch->Begin();
    }

    void SpriteRenderer::End()
    {
        this->Flush();
    }

    void SpriteRenderer::SetCamera(const std::shared_ptr<Camera2d>& cam)
    {
        this->camera = cam;
    }

    void SpriteRenderer::Flush()
    {
        batch->End();
        shader->Bind();

        if (this->camera)
        {
            glm::mat4 view = this->camera->GetViewMatrix();
            glm::mat4 proj = this->camera->GetProjectionMatrix();
            shader->set_matrix("vp", proj * view);
        }

        // bind textures
        int texIDs[16] = { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15 };
        shader->set_texture_array("tex", texIDs, 16);

        for (auto it : batch->texture_id2slot)
        {
            glActiveTexture(GL_TEXTURE0 + it.second);
            glBindTexture(GL_TEXTURE_2D, it.first);
        }

        glDrawElements(GL_TRIANGLES, batch->num_indices, GL_UNSIGNED_INT, nullptr);
        shader->Unbind();
    }
}
